#
#  checks import and export tables
#

from pykd import *
import pefile.pefile as pefile
import os.path

def checkImport( moduleName, logger ): 
           
        module = loadModule( moduleName )
        if module==None: return
        
        pe = pefile.PE( module.image(), fast_load = True )
        pe.parse_data_directories( [pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT']])
        
        for moduleEntry in pe.DIRECTORY_ENTRY_IMPORT:
            
            if moduleEntry.dll.lower() == "ntoskrnl.exe": moduleEntry.dll = "nt"
            importModule = loadModule( os.path.splitext( moduleEntry.dll )[0] )
            
            importPE = pefile.PE( importModule.image(), fast_load = True )
            importPE.parse_data_directories( [pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT']])
           
            iatEntry = module.begin() + moduleEntry.struct.FirstThunk
            
            for imp in moduleEntry.imports:

                if ptrPtr( iatEntry ) == 0:
                    break;
                
                routineOffset = 0
                
                for exp in importPE.DIRECTORY_ENTRY_EXPORT.symbols:
                    if exp.name == imp.name:
                         routineOffset = importModule.begin() + exp.address   
                         break             
                     
# TODO: forward export should be checked too                     
                if exp.forwarder != None:
                    break;                     
               
                if  routineOffset != ptrPtr( iatEntry ):
                    logger.reportTreat( hash( "%s: changed import routine %s!%s" % ( moduleName, importModule.name(), imp.name ) ),
                                      "%s: changed import routine %s!%s( %x ) to %x" % ( moduleName, importModule.name(), imp.name, routineOffset, ptrPtr( iatEntry ) ) )
                
                iatEntry += ptrSize()
                
                
def checkExport( moduleName, logger ):
    
    try:
        
        module = loadModule( moduleName )
        if module==None: return        
     
        if is64bitSystem():
            ntHeader = typedVar( "nt", "_IMAGE_NT_HEADERS64", module.begin() + ptrDWord( module.begin() + 0x3c ) )
        else:
            ntHeader = typedVar( "nt", "_IMAGE_NT_HEADERS", module.begin() + ptrDWord( module.begin() + 0x3c ) )

        if ntHeader.OptionalHeader.DataDirectory[0].Size == 0:
            return
       
        exportDirAddr = module.begin() + ntHeader.OptionalHeader.DataDirectory[0].VirtualAddress;

        namesCount = ptrDWord( exportDirAddr + 0x18 )

        funcRva = module.begin() + ptrDWord( exportDirAddr + 0x1c ) 
      
        namesRva = module.begin() + ptrDWord( exportDirAddr + 0x20 )

        ordinalRva = module.begin() + ptrDWord( exportDirAddr + 0x24 )

        for i in range( 0, namesCount ):
            exportName = loadCStr( module.begin() + ptrDWord( namesRva + 4 * i ) )
            ordinal = ptrWord( ordinalRva + 2*i )
            funcAddrByExport = module.begin() + ptrDWord( funcRva + 4 * ordinal )
            funcAddrBySymbols = module.__getattr__( exportName )        

            if  funcAddrBySymbols != 0 and funcAddrBySymbols != funcAddrByExport:
                logger.reportTreat( hash( "%s: changed kernel export routine %s" % ( moduleName, exportName ) ), 
                                          "%s change kernel export routine  %s to %#x" % (moduleName, exportName, funcAddrByExport ) )
    except:
        logger.reportWarning("%s: failed to check export dir" % moduleName )                 
                
